package com.lizemin.decorate.filter.wrapper;

import com.lizemin.decorate.util.JSON;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.util.StringUtils;
import org.springframework.web.util.HtmlUtils;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.Objects;

/**
 * @author lzm
 * @date 2025/10/6 12:25
 * @description 响应包装类, 用于对响应体中的JS脚本做转义处理
 */
public class XssResponseWrapper extends HttpServletResponseWrapper {

    private static final Logger log = LoggerFactory.getLogger(XssResponseWrapper.class);

    /**
     * 缓存响应体的输出流，用于后续过滤脚本
     */
    private final ByteArrayOutputStream cacheOutputStream = new ByteArrayOutputStream();

    private boolean useWriter = false;

    private boolean useOutputStream = false;

    /**
     * 缓存打印流
     */
    private PrintWriter cachePrintWriter;

    /**
     * 构造函数
     *
     * @param response 原始的HttpServletResponse对象
     */
    public XssResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        if (this.useWriter) {
            throw new IllegalStateException("不能同时使用getWriter和getOutputStream方法");
        }
        this.useOutputStream = true;
        return new ServletOutputStream() {
            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {

            }

            @Override
            public void write(int b) {
                // 将响应体的内容缓存起来，用于后续过滤脚本
                cacheOutputStream.write(b);
            }

            @Override
            public void write(byte[] b) throws IOException {
                cacheOutputStream.write(b);
            }

            @Override
            public void write(byte[] b, int off, int len) {
                cacheOutputStream.write(b, off, len);
            }

            @Override
            public void print(String s) throws IOException {
                cacheOutputStream.write(s.getBytes(StandardCharsets.UTF_8));
            }

            @Override
            public void println(String s) throws IOException {
                cacheOutputStream.write(s.getBytes(StandardCharsets.UTF_8));
                cacheOutputStream.write("\n".getBytes(StandardCharsets.UTF_8));
            }
        };
    }

    /**
     * 重写这个方法，是为了当调用getWriter方法时，将PrintWriter写入的内容缓存到cacheOutputStream中
     */
    @Override
    public PrintWriter getWriter() {
        if (this.useOutputStream) {
            throw new IllegalStateException("不能同时使用getWriter和getOutputStream方法");
        }
        this.useWriter = true;
        if (cachePrintWriter == null) {
            cachePrintWriter = new PrintWriter(cacheOutputStream, true, StandardCharsets.UTF_8);
        }
        return cachePrintWriter;
    }

    /**
     * 过滤响应体中的脚本
     */
    public void filterScript() throws IOException {
        if (isCommitted()) {
            log.info("响应已提交，无需过滤脚本");
            return;
        }

        // 确保所有数据都写入缓存
        if (Objects.nonNull(cachePrintWriter)) {
            cachePrintWriter.flush();
        }

        String responseBody = cacheOutputStream.toString(StandardCharsets.UTF_8);
        String contentType = super.getContentType();
        String escapedBody;
        if (StringUtils.hasText(contentType) && contentType.contains(MediaType.APPLICATION_JSON_VALUE)) {
            escapedBody = JSON.getEscapedJson(responseBody);
        } else {
            log.info("过滤前的响应体为: {}", responseBody);
            escapedBody = HtmlUtils.htmlEscape(responseBody);
            log.info("过滤后的响应体为: {}", escapedBody);
        }
        // 重置原始响应
        resetBuffer();

        // 解决浏览器的乱码问题
        setCharacterEncoding(StandardCharsets.UTF_8.name());
        setContentType("application/json;charset=UTF-8");

        // 设置内容长度
        byte[] responseBytes = escapedBody.getBytes(StandardCharsets.UTF_8);
        setContentLength(responseBytes.length);

        if (useWriter) {
            // 使用原始的打印流写入过滤后的响应内容
            try (PrintWriter originalWriter = super.getWriter()) {// 这里要注意使用这种写法关闭流，否则会导致客户端一直等待客户端的响应
                originalWriter.write(escapedBody);
                originalWriter.flush(); // 确保数据写入
            }
        } else {
            // 使用原始的输出流写入过滤后的响应内容
            try (ServletOutputStream outputStream = super.getOutputStream()) {
                outputStream.write(responseBytes);
                outputStream.flush(); // 确保数据写入
            }
        }
    }

}
